#include "PeDumpTool.h"
#include <PELib.h>
#include <qfiledialog.h>
#include <qmessagebox.h>
PeDumpTool::PeDumpTool(QWidget *parent)
	: QMainWindow(parent), dialog(NULL)
{
	ui.setupUi(this);
	QRegExp regx("[0-9]+$");
	QValidator *validator = new QIntValidator(0, 65535, this);
	QObject::connect(ui.dumpButton, SIGNAL(clicked()), this, SLOT(dumpProcess()));
	QObject::connect(ui.listhistory, SIGNAL(itemClicked(QListWidgetItem*)), this, SLOT(select_history(QListWidgetItem*)));
	QObject::connect(ui.tabWidget, SIGNAL(tabBarClicked(int)), this, SLOT(tabChange(int)));
	QObject::connect(ui.showThunk, SIGNAL(clicked()), this, SLOT(showTHUNK()));
	QObject::connect(ui.fix_int, SIGNAL(clicked()), this, SLOT(fix_int()));
	QObject::connect(ui.selectdump, SIGNAL(clicked()), this, SLOT(select_dump()));
	QObject::connect(ui.open_for_section_edit, SIGNAL(clicked()), this, SLOT(open_for_section_edit()));
	QObject::connect(ui.select2, SIGNAL(clicked()), this, SLOT(select_another_dump()));
	QObject::connect(ui.bdiff, SIGNAL(clicked()), this, SLOT(getdiff()));
	QObject::connect(ui.b_reloc, SIGNAL(clicked()), this, SLOT(create_reloc()));
	QObject::connect(ui.b_createasection, SIGNAL(clicked()), this, SLOT(create_new_Section()));
	QStringList sectionlistheader;
	ui.sectionlist->setColumnCount(6);
	ui.tabWidget->setCurrentIndex(0);
	sectionlistheader << "Section" << "RVA" << "V.SIZE" << "OFFSET" << "RAWSIZE" << "flag";
	ui.sectionlist->setHorizontalHeaderLabels(sectionlistheader);
	ui.sectionlist->horizontalHeader()->setStretchLastSection(true);
	ui.lineEdit_pid->setValidator(validator);
}
void PeDumpTool::open_for_section_edit()
{
	if (dialog == NULL)
		dialog = new QFileDialog(this);
	dialog->setAcceptMode(QFileDialog::AcceptMode::AcceptOpen);
	dialog->setFileMode(QFileDialog::FileMode::ExistingFile);
	if (dialog->exec() == QDialog::Accepted)
	{
		ui.lineEdit_filename_for_sec->setText(dialog->selectedFiles()[0]);
	}
	Current_file = dialog->selectedFiles()[0];
}

void PeDumpTool::create_new_Section()
{
	if (ui.newSecName->text().length() > 7)
	{
		QMessageBox::warning(this, "ERROR", "Name is too long");
		return;
	}
	bool ok;
	DWORD size=ui.newSecSize->text().toInt(&ok, 16);
	if (!ok)
	{
		QMessageBox::warning(this, "ERROR", "Invalid Size");
		return;
	}
	DWORD chari = ui.chari->text().toInt(&ok, 16);
	if (!ok)
	{
		QMessageBox::warning(this, "ERROR", "Invalid Size");
		return;
	}
	::addASection(ui.lineEdit_filename_for_sec->text().toStdString().c_str(),
		(ui.lineEdit_filename_for_sec->text() + "Edite").toStdString().c_str(),
		ui.newSecName->text().toStdString().c_str(), size, chari);
}
void PeDumpTool::create_reloc()
{
	std::vector<DWORD> RVA;
	RVA.clear();
	int num = ui.differences->rowCount();
	for (int i = 0; i < num; i++)
	{
		RVA.push_back(ui.differences->item(i, 0)->text().toInt(NULL,16));
	}
	::create_reloc(ui.lineEdit_dump->text().trimmed().toStdString().c_str(),RVA);
}
void PeDumpTool::getdiff()
{
	std::vector<::diff> res;
	::getPossibleReloc(ui.lineEdit_dump->text().toStdString().c_str(), ui.lineEdit_dump2->text().toStdString().c_str(), res);
	if (res.size() == 0)
	{
		Log("Getdifference failed.");
	}
	ui.differences->clear();
	int numOfDiff = res.size();
	ui.differences->setRowCount(numOfDiff);
	QStringList THUNKlistheader;
	ui.differences->setColumnCount(3);
	THUNKlistheader << "RVA" << "data file1" << "data file2";
	ui.differences->setHorizontalHeaderLabels(THUNKlistheader);
	ui.differences->horizontalHeader()->setStretchLastSection(true);
	for (int i = 0; i < numOfDiff;i++)
	{
		diff& dif = res[i];
		QTableWidgetItem* item = new QTableWidgetItem(QString::number(dif.offset, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.differences->setItem(i, 0, item);

		qlonglong  data1 = (qlonglong)dif.DataFile1;
		item = new QTableWidgetItem("0x" + QString::number(data1, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.differences->setItem(i, 1, item);

		qlonglong  data2 = (qlonglong)dif.DataFile2;
		item = new QTableWidgetItem("0x" + QString::number(data2, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.differences->setItem(i, 2, item);
	}

}
void PeDumpTool::select_dump()
{
	if (dialog == NULL)
		dialog = new QFileDialog(this);
	dialog->setAcceptMode(QFileDialog::AcceptMode::AcceptOpen);
	dialog->setFileMode(QFileDialog::FileMode::ExistingFile);
	if (dialog->exec() == QDialog::Accepted)
	{
		ui.lineEdit_dump->setText(dialog->selectedFiles()[0]);
	}

}

void PeDumpTool::select_another_dump()
{
	if (dialog == NULL)
		dialog = new QFileDialog(this);
	dialog->setAcceptMode(QFileDialog::AcceptMode::AcceptOpen);
	dialog->setFileMode(QFileDialog::FileMode::ExistingFile);
	if (dialog->exec() == QDialog::Accepted)
	{
		ui.lineEdit_dump2->setText(dialog->selectedFiles()[0]);
	}

}
void PeDumpTool::fix_int()
{
	int THUNK = ui.lineEdit_THUNK->text().toInt(NULL, 16);
	std::vector<LPVOID> functionlist;
	::getTHUNKDATA(Current_file.toStdString().c_str(), functionlist, THUNK);
	int num = functionlist.size();
	std::string dll = ui.lineEdit_DLLNAME->text().toStdString();
	if (!::validateDllName(dll.c_str(), dll.length() + 1))
		return;
	QStringList function_names=ui.functionnames->toPlainText().split("\n");
	if (function_names.length() != num)
	{
		Log("ERROR:Number of functions is incorrect!");
	}
	std::vector < std::string > functions;
	functions.clear();
	for (auto str : function_names)
		functions.push_back(str.trimmed().toStdString());
	::fix_int(Current_file.toStdString().c_str(), dll.c_str(), functions, THUNK);
}
bool PeDumpTool::showTHUNK()
{
	int THUNK = ui.lineEdit_THUNK->text().toInt(NULL,16);
	if (THUNK == 0)
		return false;
	std::vector<LPVOID> functionlist;
	::getTHUNKDATA(Current_file.toStdString().c_str(), functionlist, THUNK);
	int num = functionlist.size();
	if (num == 0)
		return false;
	ui.thunkdata_list->setRowCount(num);
	ui.thunkdata_list->setColumnCount(2);
	for (int i = 0; i <num; i++)
	{
		QTableWidgetItem* item = new QTableWidgetItem(QString::number(i,16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.thunkdata_list->setItem(i, 0, item);
		qlonglong  func = (qlonglong)functionlist[i];
		
		item = new QTableWidgetItem("0x"+QString::number(func,16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.thunkdata_list->setItem(i, 1, item);
	}
	QStringList THUNKlistheader;
	ui.thunkdata_list->setColumnCount(2);
	THUNKlistheader << "No." << "Funtcion_Address";
	ui.thunkdata_list->setHorizontalHeaderLabels(THUNKlistheader);
	ui.thunkdata_list->horizontalHeader()->setStretchLastSection(true);
	return true;
}
void PeDumpTool::tabChange(int index)
{
	switch (index)
	{
	case 1:
	{
		showSection(Current_file);
		break;
	}
	}
}

void PeDumpTool::SectionInfoChanged(QTableWidgetItem* item)
{
	IMAGE_SECTION_HEADER header;
	
	int row=item->row();
	strcpy_s((char*)header.Name, 8,ui.sectionlist->item(row, 0)->text().toStdString().c_str());
	header.VirtualAddress = +ui.sectionlist->item(row, 1)->text().toLongLong(NULL, 16);
	header.Misc.VirtualSize = ui.sectionlist->item(row, 2)->text().toLongLong(NULL, 16);
	header.PointerToRawData = ui.sectionlist->item(row, 3)->text().toLongLong(NULL, 16);
	header.SizeOfRawData = ui.sectionlist->item(row, 4)->text().toLongLong(NULL, 16);
	header.Characteristics = ui.sectionlist->item(row, 5)->text().toLongLong(NULL, 16);
	SetSectionHeader(Current_file.toStdString().c_str(), row, header);
}
void PeDumpTool::showSection(QString filename)
{
	QObject::disconnect(ui.sectionlist, SIGNAL(itemChanged(QTableWidgetItem*)), this, SLOT(SectionInfoChanged(QTableWidgetItem*)));
	std::vector<IMAGE_SECTION_HEADER> sectioninfo;
	::getSectionList(filename.toStdString().c_str(), sectioninfo);
	if (sectioninfo.size() == 0)
		return;
	ui.sectionlist->clear();
	int numOfSection = sectioninfo.size();
	ui.sectionlist->setRowCount(numOfSection);
	
	for (int i = 0; i < numOfSection;i++)
	{
		QTableWidgetItem* item=new QTableWidgetItem((char*)sectioninfo[i].Name);
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.sectionlist->setItem(i, 0, item);
		item = new QTableWidgetItem(QString::number(sectioninfo[i].VirtualAddress, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.sectionlist->setItem(i, 1, item);
		item = new QTableWidgetItem(QString::number(sectioninfo[i].Misc.VirtualSize, 16));
		ui.sectionlist->setItem(i, 2, item);
		item = new QTableWidgetItem(QString::number(sectioninfo[i].PointerToRawData, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.sectionlist->setItem(i, 3, item);
		item = new QTableWidgetItem(QString::number(sectioninfo[i].SizeOfRawData, 16));
		item->setFlags(item->flags() & ~Qt::ItemIsEditable);
		ui.sectionlist->setItem(i, 4, item);
		item = new  QTableWidgetItem(QString::number(sectioninfo[i].Characteristics, 16));
		ui.sectionlist->setItem(i, 5, item);
	}
	QStringList sectionlistheader;
	ui.sectionlist->setColumnCount(6);
	sectionlistheader << "Section" << "RVA" << "V.SIZE" << "OFFSET" << "RAWSIZE" << "flag";
	ui.sectionlist->setHorizontalHeaderLabels(sectionlistheader);
	ui.sectionlist->horizontalHeader()->setStretchLastSection(true);
	QObject::connect(ui.sectionlist, SIGNAL(itemChanged(QTableWidgetItem*)), this, SLOT(SectionInfoChanged(QTableWidgetItem*)));

}
void PeDumpTool::Log(const QString& str)
{
	ui.OutputL->append(str);
}
void PeDumpTool::select_history(QListWidgetItem *item)
{
	Current_file = item->text();
	Log("cur:"+Current_file);
}
bool PeDumpTool::dumpProcess()
{
	int pid = ui.lineEdit_pid->text().toInt();

	QString outName = ui.lineEdit_NewFIleName->text().toStdString().data();
	QString moduleName = ui.lineEdit_modulename->text().toStdString().data();
	QString oep = ui.lineEdit_oep->text().toStdString().data();
	DWORD oepnum;
	if (oep != QString()&&sscanf(oep.toStdString().c_str(), "%x", &oepnum) < 1)
	{
		ui.OutputL->append("oep invalid\n");
		return false;
	}
	std::string res = ::dumpPe(pid, moduleName.toStdString().c_str(), outName.toStdString().c_str(), oepnum);
	if (res == std::string())
	{
		ui.OutputL->append("Dump Failed\n");
		return false;
	}
	ui.lineEdit_NewFIleName->setText(res.c_str());
	QStringList history;
	QString log = "Dump " + ui.lineEdit_pid->text() + " to " + res.c_str()+"\n";
	Log(log);
	dumpedhistory.push_back(res);
	ui.listhistory->addItem(res.c_str());
	return true;
}